from .openweb import load_context
class ContextLoader:
    def __init__(self, collected_filenames):
        self.collected_filenames = collected_filenames
        self.filename_to_path = {}

    def load_files(self):
        # 遍历 collected_filenames 并填充字典
        for file_info_list in self.collected_filenames:
            for file_info in file_info_list:
                self.filename_to_path[file_info['filename']] = file_info['full_path']

    def get_file_paths(self):
        # 定义常量
        COMMUNITY_REPORT_TABLE = "create_final_community_reports.parquet"
        ENTITY_TABLE = "create_final_nodes.parquet"
        ENTITY_EMBEDDING_TABLE = "create_final_entities.parquet"
        RELATIONSHIP_TABLE = "create_final_relationships.parquet"
        TEXT_UNIT_TABLE = "create_final_text_units.parquet"

        # 获取文件路径
        community_report_table_path = self.filename_to_path.get(COMMUNITY_REPORT_TABLE)
        entity_table_path = self.filename_to_path.get(ENTITY_TABLE)
        entity_embedding_table_path = self.filename_to_path.get(ENTITY_EMBEDDING_TABLE)
        relationship_table_path = self.filename_to_path.get(RELATIONSHIP_TABLE)
        text_unit_table_path = self.filename_to_path.get(TEXT_UNIT_TABLE)

        return {
            "community_report_table_path": community_report_table_path,
            "entity_table_path": entity_table_path,
            "entity_embedding_table_path": entity_embedding_table_path,
            "relationship_table_path": relationship_table_path,
            "text_unit_table_path": text_unit_table_path,
        }

    def load_context(self):
        paths = self.get_file_paths()

        # 使用获取的路径加载上下文
        load_context(
            entity_table_path=str(paths["entity_table_path"]),
            entity_embedding_table_path=str(paths["entity_embedding_table_path"]),
            relationship_table_path=str(paths["relationship_table_path"]),
            community_report_table_path=str(paths["community_report_table_path"]),
            text_unit_table_path=str(paths["text_unit_table_path"])
        )
